# util/calculation.py

import numpy as np
import time
from scipy.special import jv, ellipk, ellipe
from scipy.integrate import trapezoid, quad
from joblib import Parallel, delayed
import itertools


# ===================================================================
# ---         Default Parameters & Configuration                  ---
# ===================================================================

def get_default_wpt_parameters():
    """Returns a dictionary with default physical parameters for the WPT system."""
    return {
        'omega': 2 * np.pi * 85000, 'Us': 50.0, 'mu0': 4 * np.pi * 1e-7,
        'kc': 0.175 / (2 * np.pi * 100000), 'n0': 300, 'rliz': 0.05e-3,
        'sigmac': 58e6, 'I': 20.25e-3, 'rc': 0.87e-3,
        'muc': 0.999991 * 4 * np.pi * 1e-7, 'kelse': 0.9, 'h1': 0.002,
        'femiur': 3300.0, 'fesigma': 1 / 6.0, 'h': 50e-3,
        'xs': 0.0, 'RL': 5.0,
    }


# ===================================================================
# ---         1. Resistance Calculation (Parallelized)            ---
# ===================================================================

def calculate_resistance(params: dict):
    """
    Calculates the total AC resistance for Tx and Rx coils using parallel processing.
    """
    Np, Ns = params['Np'], params['Ns']

    Rp_turns = Parallel(n_jobs=1)(delayed(_calculate_single_turn_resistance)(i, 'tx', params) for i in range(Np))
    Rs_turns = Parallel(n_jobs=1)(delayed(_calculate_single_turn_resistance)(i, 'rx', params) for i in range(Ns))

    return sum(Rp_turns), sum(Rs_turns)


def _calculate_single_turn_resistance(i, coil_type, params):
    """
    Calculates the complete AC resistance (conduction + proximity) for a single turn.
    This function is designed to be called in parallel by joblib.
    """
    n0, rliz, sigmac = params['n0'], params['rliz'], params['sigmac']
    kc, omega = params['kc'], params['omega']
    rad_array = params['radtx'] if coil_type == 'tx' else params['radrx']
    radius = rad_array[i]

    # Conduction (DC) Resistance, adjusted for skin effect
    R_cond = (2 * np.pi * radius) / (n0 * np.pi * rliz ** 2 * sigmac)
    R_cond_ac = R_cond * (1 + kc * omega) * 1.05

    # Proximity Effect Resistance
    R_prox = _calculate_proximity_resistance_for_turn(i, coil_type, params)

    return R_cond_ac + R_prox


def _calculate_proximity_resistance_for_turn(current_turn, coil_type, params):
    """Helper to calculate proximity resistance, which depends on the H-field."""
    I, rc, n0, muc, sigmac, rliz, omega = [params[k] for k in ['I', 'rc', 'n0', 'muc', 'sigmac', 'rliz', 'omega']]
    rad_array = params['radtx'] if coil_type == 'tx' else params['radrx']
    radius = rad_array[current_turn]
    Havg_squared = _calculate_havg_squared(current_turn, coil_type, params)
    return (n0 * np.pi ** 2 * muc ** 2 * sigmac * rliz ** 4 * radius * Havg_squared * omega ** 2) / (2 * I ** 2)


def _calculate_havg_squared(obs_turn_idx, coil_type_obs, params):
    """Calculates the average squared magnetic field over a turn's cross-section."""
    I, rc, kelse, h1, Np, Ns, radtx, radrx = [params[k] for k in ['I', 'rc', 'kelse', 'h1', 'Np', 'Ns', 'radtx', 'radrx']]
    A, hs = 3.0 / 8.0, 7.0 * rc / 24.0
    Wn, theta_points, rho_points = np.array([1, 3, 3, 1]), np.pi * np.array([.25, .5, .75, 1.25, 1.5, 1.75]), rc * np.array(
        [1 / 8, 10 / 24, 17 / 24, 1])
    H_tempout = 0.0
    obs_rad_array = radtx if coil_type_obs == 'tx' else radrx
    obs_radius = obs_rad_array[obs_turn_idx]
    source_rad_array = radtx if coil_type_obs == 'tx' else radrx
    N_source = Np if coil_type_obs == 'tx' else Ns

    for theta in theta_points:
        for n, rho in enumerate(rho_points):
            r_obs, z_obs = obs_radius + rho * np.cos(theta), rho * np.sin(theta)
            Hr_total, Hz_total = 0.0, 0.0
            for p in range(N_source):
                if p == obs_turn_idx:
                    Hiji = I * rho / (2 * np.pi * rc ** 2)
                    Hr_p, Hz_p = Hiji * np.sin(theta), Hiji * np.cos(theta)
                else:
                    Hr_p, Hz_p = _get_h_components_from_turn(r_obs, z_obs, source_rad_array[p], coil_type_obs, params)
                    Hr_p *= kelse
                    Hz_p *= kelse
                Hr_total += Hr_p
                Hz_total += Hz_p
            H_tempout += Wn[n] * (Hr_total ** 2 + Hz_total ** 2) * rho

    return H_tempout * A * hs / (3 * rc ** 2)


def _get_h_components_from_turn(r_obs, z_obs, source_radius, coil_type_source, params):
    """Calculates H-field from a single source turn filament."""
    I, rc, h1 = params['I'], params['rc'], params['h1']
    a_values = source_radius + np.array([-rc / 2, -rc / 4, 0, rc / 4, rc / 2])
    Hr_vals, Hz_vals = np.zeros_like(a_values), np.zeros_like(a_values)

    for i, a in enumerate(a_values):
        if r_obs < 1e-9: continue
        m1_sq = np.clip(4 * a * r_obs / ((a + r_obs) ** 2 + z_obs ** 2), 0, 1 - 1e-9)
        K1, E1 = ellipk(m1_sq), ellipe(m1_sq)
        c1 = I / (np.pi * np.sqrt((a + r_obs) ** 2 + z_obs ** 2))
        hr1 = c1 * (z_obs / (2 * r_obs)) * (((a ** 2 + r_obs ** 2 + z_obs ** 2) / ((a - r_obs) ** 2 + z_obs ** 2)) * E1 - K1)
        hz1 = (c1 / 2) * (K1 + ((a ** 2 - r_obs ** 2 - z_obs ** 2) / ((a - r_obs) ** 2 + z_obs ** 2)) * E1)

        if coil_type_source == 'tx':
            m2_sq = np.clip(4 * a * r_obs / ((a + r_obs) ** 2 + (z_obs + 2 * h1) ** 2), 0, 1 - 1e-9)
            K2, E2 = ellipk(m2_sq), ellipe(m2_sq)
            c2 = I / (np.pi * np.sqrt((a + r_obs) ** 2 + (z_obs + 2 * h1) ** 2))
            hr2 = c2 * ((z_obs + 2 * h1) / (2 * r_obs)) * (
                        ((a ** 2 + r_obs ** 2 + (z_obs + 2 * h1) ** 2) / ((a - r_obs) ** 2 + (z_obs + 2 * h1) ** 2)) * E2 - K2)
            hz2 = (c2 / 2) * (K2 + ((a ** 2 - r_obs ** 2 - (z_obs + 2 * h1) ** 2) / ((a - r_obs) ** 2 + (z_obs + 2 * h1) ** 2)) * E2)
            Hr_vals[i], Hz_vals[i] = hr1 + hr2, hz1 + hz2
        else:
            Hr_vals[i], Hz_vals[i] = hr1, hz1

    return np.mean(Hr_vals), np.mean(Hz_vals)


# ===================================================================
# ---         2. Mutual Inductance Calculation (Vectorized)       ---
# ===================================================================

def calculate_mutual_inductance(params: dict):
    """Calculates mutual inductance using a high-speed vectorized method."""
    Np, Ns, radtx, radrx = params['Np'], params['Ns'], params['radtx'], params['radrx']
    omega, mu0, h, xs = params['omega'], params['mu0'], params['h'], params['xs']
    h1, femiur, fesigma, df = params['h1'], params['femiur'], params['fesigma'], params['df']

    if Np == 0 or Ns == 0: return 0.0

    K_POINTS, K_MAX = 2000, 1000.0
    k_grid = np.linspace(1e-6, K_MAX, K_POINTS)
    radtx_b, radrx_b = radtx[:, np.newaxis], radrx[:, np.newaxis]
    bessel_prod = jv(1, radtx_b * k_grid)[:, np.newaxis, :] * jv(1, radrx_b * k_grid)[np.newaxis, :, :]
    bessel_j0 = jv(0, xs * k_grid)

    # Air inductance part
    exp_term = np.exp(-k_grid * h)
    integrand_Mair = bessel_prod * bessel_j0 * exp_term
    M_air_matrix = trapezoid(integrand_Mair, k_grid, axis=2)
    M_air = np.pi * mu0 * np.sum(radtx_b * radrx_b.T * M_air_matrix)

    # Ferrite inductance part
    eta_k = np.sqrt(k_grid ** 2 + 1j * omega * femiur * mu0 * fesigma)
    faik_val = np.divide(femiur - eta_k / k_grid, femiur + eta_k / k_grid, out=np.zeros_like(k_grid, dtype=complex), where=k_grid != 0)
    lambda_val = faik_val * (1 - np.exp(-2 * eta_k * df)) / (1 - faik_val ** 2 * np.exp(-2 * eta_k * df))
    exp_term_ferrite = np.exp(-k_grid * (h + 2 * h1))
    integrand_Zf = bessel_prod * bessel_j0 * lambda_val * exp_term_ferrite
    Zf_matrix = trapezoid(integrand_Zf, k_grid, axis=2)
    Z_ferrite = 1j * omega * mu0 * np.pi * np.sum(radtx_b * radrx_b.T * Zf_matrix)

    M_ferrite = np.imag(Z_ferrite) / omega if omega > 0 else 0.0
    return M_ferrite + M_air


# ===================================================================
# ---         3. Transmission Characteristics Calculation         ---
# ===================================================================

def calculate_transmission_characteristics(M, Rp, Rs, params: dict):
    """Calculates system efficiency and output power."""
    omega, Us, RL = params['omega'], params['Us'], params['RL']

    denominator_eff = Rp * (Rs + RL) ** 2 + omega ** 2 * M ** 2 * (Rs + RL)
    efficiency = (omega ** 2 * M ** 2 * RL) / denominator_eff if denominator_eff != 0 else 0

    denominator_pout = (Rp * (Rs + RL) + omega ** 2 * M ** 2) ** 2
    Pout = (omega ** 2 * M ** 2 * Us ** 2 * RL) / denominator_pout if denominator_pout != 0 else 0

    return efficiency, Pout


# ===================================================================
# ---         4. Main All-in-One Function                         ---
# ===================================================================

def calculate_wpt_performance(coil_params: dict):
    """
    A comprehensive function that runs all calculations in sequence.
    Ideal for single-shot analysis, but not for use inside the RL loop.
    """
    # 1. Prepare parameters
    params = get_default_wpt_parameters()
    params.update(coil_params)

    Nin, Nout, Ns = params['Nin'], params['Nout'], params['Ns']
    rin, din, rout, dout = params['rin'], params['din'], params['rout'], params['dout']
    rs, ds = params['rs'], params['ds']

    Np = Nin + Nout
    params['Np'], params['Ns'] = Np, Ns

    radtx = np.zeros(Np)
    radtx[:Nin] = rin + np.arange(Nin) * din
    radtx[Nin:] = rout + np.arange(Nout) * dout
    params['radtx'] = radtx
    params['radrx'] = rs + np.arange(Ns) * ds

    # 2. Run calculations sequentially
    Rp_val, Rs_val = calculate_resistance(params)
    M_val = calculate_mutual_inductance(params)
    eta_val, P_load_val = calculate_transmission_characteristics(M_val, Rp_val, Rs_val, params)

    return {
        "Rp": Rp_val,
        "Rs": Rs_val,
        "M": M_val,
        "efficiency": eta_val,
        "P_load": P_load_val,
    }


# ===================================================================
# ---                  Module Test Block                          ---
# ===================================================================

if __name__ == '__main__':
    # --- Example: How to use the modular functions ---

    my_coil_parameters = {'h': 0.05, 'xs': 0.0, 'RL': 5.31035019, 'Nin': 9, 'rin': 0.18125987, 'din': 0.00503551, 'Nout': 11
        , 'rout': 0.25138447, 'dout': 0.00556386,
                          'Ns': 13, 'rs': 0.08248829, 'ds': 0.00409771, 'df': 0.00900000}

    print("--- Testing the comprehensive 'calculate_wpt_performance' function ---")
    start_time = time.time()
    results = calculate_wpt_performance(my_coil_parameters)
    end_time = time.time()
    print(f"Calculation finished in {end_time - start_time:.4f} seconds.")

    if results:
        print("\n--- WPT System Performance Results ---")
        print(f"Transmitter Coil Resistance (Rp): {results['Rp']:.6f} Ω")
        print(f"Receiver Coil Resistance (Rs):  {results['Rs']:.6f} Ω")
        print(f"Mutual Inductance (M):          {results['M'] * 1e6:.6f} µH")
        print(f"Output Power (P_load):          {results['P_load']:.6f} W")
        print(f"Transmission Efficiency (eta):  {results['efficiency'] * 100:.4f} %")
        print("--------------------------------------")

